package interviewpreparationkit.dictionary


/**
 * [Count Triplets](https://www.hackerrank.com/challenges/count-triplets-1)
 */
// Complete the countTriplets function below.
fun countTriplets(arr: Array<Long>, r: Long): Long {
    if (r == 1L) {
        val map = mutableMapOf<Long, Long>()
        arr.forEach { map[it] = map.getOrDefault(it, 0L) + 1L }
        return map.values.filter { it > 2L }.map { it * (it - 1L) * (it - 2L) / 6L }.sum()
    }

    val map1 = mutableMapOf<Long, Long>()
    val map2 = mutableMapOf<Long, Long>()
    var count: Long = 0

    for (i in arr.lastIndex downTo 0) {
        val it = arr[i]
        val itR = it * r
        if (map1.containsKey(itR)) {
            val c = map1[itR]
            map2[it] = map2.getOrDefault(it, 0L) + c!!
        }

        if (map2.containsKey(itR))
            count += map2[itR]!!

        map1[it] = map1.getOrDefault(it, 0L) + 1L
    }

    return count
}